from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from flax import linen as nn

from jax import tree_util

from cca_zoo.linear import CCA, rCCA
from cca_zoo.model_selection import GridSearchCV

import seaborn as sns
import logging
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd

import numpy as np
from scipy.linalg import qr, svd, inv
import logging

from scipy.spatial.distance import pdist, squareform
from scipy.stats import pearsonr

from bokeh.plotting import figure, show, output_file
from bokeh.models import ColumnDataSource, HoverTool
from bokeh.palettes import Category20

from joblib import dump


def vector_projection(A, X):
    """
    计算向量 A 在基 X 上的投影。

    参数：
    A: numpy.ndarray，形状为 (128,)
    X: numpy.ndarray，形状为 (4, 128)，线性无关的基向量 X0, X1, X2, X3

    返回：
    projections: numpy.ndarray，形状为 (4,)，A 在 X0, X1, X2, X3 上的投影值
    """
        # 使用 QR 分解将 X 正交化
    Q, R = np.linalg.qr(X.T)
    # Q 的形状为 (128, 4)，正交列向量
    U = Q[:, :4].T  # 形状转换为 (4, 128)

    # 计算 A 在正交基上的当前投影
    current_projections = U @ A  # 形状为 (4,)

    return current_projections

def adjust_vector_projection_orthonormal(A, X, desired_projections):
    """
    调整向量 A，使其在正交基 X 上的投影等于 desired_projections。

    参数：
    A: numpy.ndarray，形状为 (128,)
    X: numpy.ndarray，形状为 (4, 128)，线性无关的基向量 X0, X1, X2, X3
    desired_projections: numpy.ndarray，形状为 (4,)，目标投影值 y0, y1, y2, y3

    返回：
    A_prime: numpy.ndarray，形状为 (128,)，调整后的向量
    """
    # 使用 QR 分解将 X 正交化
    Q, R = np.linalg.qr(X.T)
    # Q 的形状为 (128, 4)，正交列向量
    U = Q[:, :4].T  # 形状转换为 (4, 128)

    # 计算 A 在正交基上的当前投影
    current_projections = U @ A  # 形状为 (4,)

    # print(current_projections)

    # 计算需要调整的投影差值
    delta = desired_projections - current_projections  # 形状为 (4,)

    # 调整向量 A
    A_prime = A + U.T @ delta  # U.T 形状为 (128,4), delta 是 (4,)

    return A_prime


def transform_cca(X, A):

    return X @ A

def inverse_cca(A, U, data_prop):
    """
    Inverse CCA: Given the canonical coefficients `A`, the canonical scores `U`, the standard deviation `std_x` and the mean `mean_x` of the original data, return the original data `X`
    """

    # check if A is a square matrix
    assert A.shape[0] == A.shape[1]

    A_inverse = np.linalg.pinv(A)
    _packed_rnn_states = U @ A_inverse
    # _packed_rnn_states = _packed_rnn_states * data_prop[2] + data_prop[0]
    return _packed_rnn_states


def canoncorr(X0: np.array, Y0: np.array, fullReturn: bool = False) -> np.array:
    """
    Canonical Correlation Analysis (CCA) with added diagnostics and preprocessing
    
    Parameters:
    X, Y: (samples/observations) x (features) matrix
    fullReturn: whether all outputs should be returned or just `r` be returned
    
    Returns: A, B, r, U, V (if fullReturn is True) or just r (if fullReturn is False)
    A, B: Canonical coefficients for X and Y
    U, V: Canonical scores for the variables X and Y
    r: Canonical correlations
    """
    n, p1 = X0.shape
    p2 = Y0.shape[1]
    
    # # Data diagnostics
    # print(f"X shape: {X0.shape}, Y shape: {Y0.shape}")
    # print(f"X condition number: {np.linalg.cond(X0)}")
    # print(f"Y condition number: {np.linalg.cond(Y0)}")
    
    if p1 >= n or p2 >= n:
        logging.warning('Not enough samples, might cause problems')

    # Preprocessing: Standardize the variables
    x_mean, y_mean = np.mean(X0, 0), np.mean(Y0, 0)
    x_std, y_std = np.std(X0, 0), np.std(Y0, 0)
    X = (X0 - x_mean) / x_std
    Y = (Y0 - y_mean) / y_std

    # print(np.std(X0, 0), np.std(Y0, 0))

    # Factor the inputs, and find a full rank set of columns if necessary
    Q1, T11, perm1 = qr(X, mode='economic', pivoting=True)
    Q2, T22, perm2 = qr(Y, mode='economic', pivoting=True)

    # Determine ranks with a more stringent threshold
    tol = np.finfo(float).eps * 100  # Increased tolerance
    rankX = np.sum(np.abs(np.diagonal(T11)) > tol * np.abs(T11[0, 0]))
    rankY = np.sum(np.abs(np.diagonal(T22)) > tol * np.abs(T22[0, 0]))

    # print(f"Rank of X: {rankX}, Rank of Y: {rankY}")

    if rankX == 0:
        raise ValueError('X has zero rank')
    elif rankX < p1:
        logging.warning('X is not full rank')
        Q1 = Q1[:, :rankX]
        T11 = T11[:rankX, :rankX]

    if rankY == 0:
        raise ValueError('Y has zero rank')
    elif rankY < p2:
        logging.warning('Y is not full rank')
        Q2 = Q2[:, :rankY]
        T22 = T22[:rankY, :rankY]

    # Compute canonical coefficients and canonical correlations
    d = min(rankX, rankY)
    # L, D, M = svd(Q1.T @ Q2, full_matrices=False)
    L,D,M = svd(Q1.T @ Q2, full_matrices=True, check_finite=True, lapack_driver='gesdd')
    M = M.T

    A = inv(T11) @ L[:, :d] * np.sqrt(n - 1)
    B = inv(T22) @ M[:, :d] * np.sqrt(n - 1)
    r = D[:d]
    
    # Remove roundoff errors
    r = np.clip(r, 0, 1)

    if not fullReturn:
        return r

    # Put coefficients back to their full size and correct order
    A_full = np.zeros((p1, d))
    A_full[perm1, :] = np.vstack((A, np.zeros((p1 - rankX, d))))
    
    B_full = np.zeros((p2, d))
    B_full[perm2, :] = np.vstack((B, np.zeros((p2 - rankY, d))))

    # Compute the canonical variates
    U = X @ A_full
    V = Y @ B_full

    return A_full, B_full, r, U, V, [x_mean, y_mean, x_std, y_std]


def plot_lollipop(scores, title="Lollipop Plot of Scores", xlabel="Index", ylabel="Score", figsize=(12, 6)):
    """
    绘制带连线和数值标注的棒棒糖图，数值保留两位小数，x轴标签从1开始
    
    参数:
    scores : list 或 numpy array
        要绘制的数据
    title : str, 可选
        图表标题
    xlabel : str, 可选
        x轴标签
    ylabel : str, 可选
        y轴标签
    figsize : tuple, 可选
        图表大小
    """
    
    # 创建对应的 x 值
    x = np.arange(len(scores))

    # 创建图表
    fig, ax = plt.subplots(figsize=figsize)

    # 绘制垂直线
    ax.vlines(x, 0, scores, colors='gray', lw=1, alpha=0.5)

    # 绘制数据点并连线
    ax.plot(x, scores, color='black', marker='o', markersize=8, linestyle='-', linewidth=1)

    # 添加数值标注，保留两位小数
    for i, score in enumerate(scores):
        ax.annotate(f'{score:.2f}', (i, score), textcoords="offset points", 
                    xytext=(0,10), ha='center', va='bottom', fontsize=8)

    # 设置坐标轴范围
    ax.set_xlim(-0.5, len(scores) - 0.5)
    ax.set_ylim(0, max(scores) * 1.2)  # 增加上界以容纳标注

    # 设置标题和标签
    # ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)

    # 移除顶部和右侧边框
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    # 移除网格线
    ax.grid(False)

    # # 调整 x 轴刻度，标签从1开始
    # ax.set_xticks(x)
    # ax.set_xticklabels(range(1, (len(scores) + 1)))

    # # 调整 x 轴刻度，每个数字为横坐标的10倍
    # ax.set_xticks(x)
    # ax.set_xticklabels([str(i*10) for i in range(len(scores))])

    # 设置 y 范围 0-1
    ax.set_ylim(0, 1.1)

    plt.tight_layout()
    plt.show()


def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@partial(jax.jit, static_argnums=(3,))
def model_forward(variables, state, x, model):
    """ forward pass of the model
    """
    return model.apply(variables, state, x)

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# load landscape and states from file
def load_task(pth = "./logs/task.json", display = True):
    # open json file
    with open(pth, "r") as f:
        data = json.load(f)
        landscape = data["data"]
        state = data["state"]
        goal = data["goal"]
        if display:
            print("state: ", state)
            print("goal: ", goal)
            print("landscape: ", landscape)
    return landscape, state, goal


def preprocess(trjs, max_length):
    processed_trjs = []
    for i in range(trjs.shape[0]):
        trj = trjs[i]
        if trj.shape[0] < max_length:
            last_element = trj[-1]
            processed_trj = np.concatenate([trj, np.repeat(last_element[np.newaxis,:], max_length - trj.shape[0], axis=0)], axis=0)
            processed_trjs.append(processed_trj)
        else:
            processed_trjs.append(trj)
    processed_trjs = np.array(processed_trjs)
    return processed_trjs
    
@jax.jit
def build_radiance_field(p):
    img = jnp.zeros((21, 21))
    top = 10
    bottom = 0
    effective_radius = 21*1.414
    px, py = p[0], p[1]

    x_coord_map = jnp.arange(img.shape[0])
    x_coord_map = jnp.repeat(x_coord_map[:, jnp.newaxis], img.shape[1], axis=1)

    y_coord_map = jnp.arange(img.shape[1])
    y_coord_map = jnp.repeat(y_coord_map[jnp.newaxis, :], img.shape[0], axis=0)
    
    dist = jnp.sqrt((x_coord_map - px)**2 + (y_coord_map - py)**2)
    img = jnp.where(dist < effective_radius,
                    (top - bottom) * (effective_radius - dist) / effective_radius + bottom,
                    img)
    return img

@jax.jit
def get_max_radiance_field(imgs):
    img = jnp.max(imgs, axis=0)
    return img

@jax.jit
def build_ridge(A):

    # 将 A 的第一个点对齐到 (10,10)
    A = A-A[0]+jnp.array([10,10])

    # 使用 jnp.map 来并行计算每个辐射场图像
    imgs = jax.vmap(build_radiance_field)(A)
    img = get_max_radiance_field(imgs)
    return img

build_ridge_vmap = jax.vmap(build_ridge)


if __name__ == "__main__":

    # 载入三个数组
    intr_cca_array = np.load("./logs/intr_cca_array.npy")
    transformd_cca_representations_array = np.load("./logs/transformd_cca_representations_array.npy")
    stable_trj_center_cca_array = np.load("./logs/stable_trj_center_cca_array.npy")

    print("shape of intr_cca_array: ", intr_cca_array.shape)
    print("shape of transformd_cca_representations_array: ", transformd_cca_representations_array.shape)
    print("shape of stable_trj_center_cca_array: ", stable_trj_center_cca_array.shape)

    num_envs = intr_cca_array.shape[0]

    vectors_avg_full = []

    for i in range(num_envs):
        progress_bar(i, num_envs)

        # 计算当前环境的向量场的平均向量
        intr_cca = intr_cca_array[i]
        avg_intr_cca = np.mean(intr_cca, axis=0)

        stable_trj_center = stable_trj_center_cca_array[i][0]

        vectors_avg_full.append(avg_intr_cca)

    vectors_avg_full = np.array(vectors_avg_full)

    # 去除 stable_trj_center_cca_array 中间的维度
    stable_trj_center_cca_array = np.squeeze(stable_trj_center_cca_array, axis=1)

    print("shape of vectors_avg_full: ", vectors_avg_full.shape)
    print("shape of stable_trj_center_cca_array: ", stable_trj_center_cca_array.shape)

    # 作为对照，生成一个 stable_trj_center_cca_array 形状一模一样的，值分布在 [-1,1] 之间的随机数组

    dim_limit = 20
    train_view_1 = vectors_avg_full[:,0:dim_limit]
    train_view_2 = stable_trj_center_cca_array[:,0:dim_limit]

    random_array = np.random.uniform(-1, 1, train_view_2.shape)
    train_view_2_random = random_array

    # 对 train_view_1/2 进行 pca
    pca = PCA()
    pca.fit(train_view_1)
    train_view_1 = pca.transform(train_view_1)
    pca1 = PCA()
    pca1.fit(train_view_2)
    train_view_2 = pca1.transform(train_view_2)
    pca2 = PCA()
    pca2.fit(train_view_2_random)
    train_view_2_random = pca2.transform(train_view_2_random)

    # 保存 PCA 模型
    dump(pca, './logs/pca_model.pkl')
    dump(pca1, './logs/pca1_model.pkl')

    # using CCA from article 'Preserved neural dynamics across animals performing similar behaviour'
    A, B, r, U, V, data_prop = canoncorr(train_view_1, train_view_2, fullReturn=True)
    print(A.shape, B.shape, r.shape, U.shape, V.shape)

    # 保存 A 和 B 到文件
    np.save("./logs/A_align.npy", A)
    np.save("./logs/B_align.npy", B)

    # 计算一个“正交性”矩阵，用于测量A中的每个向量之间的正交性，其尺寸为 A 的行数 x A 的行数
    orthogonality_matrix = np.abs(np.dot(A, A.T))
    # 将这个矩阵绘制出来
    plt.imshow(orthogonality_matrix, cmap='hot', interpolation='nearest')
    plt.colorbar()
    plt.title("Orthogonality Matrix of A")
    plt.show()
    

    # 绘制 CCA 的结果
    plot_lollipop(r, title="Canonical Correlation Coefficients", xlabel="Canonical Component", ylabel="Correlation Coefficient")

    A1, B1, r1, U1, V1, data_prop1 = canoncorr(train_view_1, train_view_2_random, fullReturn=True)
    print(A1.shape, B1.shape, r1.shape, U1.shape, V1.shape)

    # 绘制 CCA 的结果
    plot_lollipop(r1, title="Canonical Correlation Coefficients", xlabel="Canonical Component", ylabel="Correlation Coefficient")

    # 计算cca变换之后每一个维度的 variance
    dims = U.shape[1]
    A_t = A.T
    # 将 train_view_1 通过 A_t 变换到新的空间
    transformed_train_view_1 = train_view_1 @ A_t
    print("shape of transformed_train_view_1: ", transformed_train_view_1.shape)
    transformed_train_view_1_t = transformed_train_view_1.T
    # 计算每个维度的 variance
    variances = np.var(transformed_train_view_1, axis=0)

    # 绘制每个维度的 variance
    plt.plot(variances)
    plt.title("Variance of each dimension after CCA")
    plt.show()
    

    # for i in range(10):

    #     cm1, cm2, cm3 = i * 3, i * 3 + 1, i * 3 + 2

    #     # 将 U 和 V 绘制到两个并排的 3D 图中，并且通过实现鼠标交互函数来同步两个图的旋转
    #     fig = plt.figure()
    #     ax1 = fig.add_subplot(121, projection='3d')
    #     ax2 = fig.add_subplot(122, projection='3d')

    #     # 定义鼠标事件处理函数
    #     def on_mouse_event(event):
    #         if event.inaxes == ax1:
    #             # 如果事件发生在ax1中，更新ax2的视角
    #             ax2.view_init(elev=ax1.elev, azim=ax1.azim)
    #         elif event.inaxes == ax2:
    #             # 如果事件发生在ax2中，更新ax1的视角
    #             ax1.view_init(elev=ax2.elev, azim=ax2.azim)
    #         fig.canvas.draw()
    #     # 绑定鼠标事件处理函数
    #     fig.canvas.mpl_connect('motion_notify_event', on_mouse_event)

    #     for j in range(10):
    #         start = j * (U.shape[0]//10)
    #         end = (j+1) * (U.shape[0]//10)
    #         ax1.scatter(U[start:end, cm1], U[start:end, cm2], U[start:end, cm3], c='b', marker='o')
    #         ax2.scatter(V[start:end, cm1], V[start:end, cm2], V[start:end, cm3], c='r', marker='o')

    #     # ax1 的标题为 neural dynamics
    #     ax1.set_title('Neural Dynamics')
    #     # ax2 的标题为 solution structure
    #     ax2.set_title('Solution Structure')

    #     ax1.set_xlabel(f'Canonical mode {cm1}')
    #     ax1.set_ylabel(f'Canonical mode {cm2}')
    #     ax1.set_zlabel(f'Canonical mode {cm3}')
    #     ax2.set_xlabel(f'Canonical mode {cm1}')
    #     ax2.set_ylabel(f'Canonical mode {cm2}')
    #     ax2.set_zlabel(f'Canonical mode {cm3}')

    #     # 设置两个 view 为 top-down 视角
    #     ax1.view_init(elev=90, azim=0)
    #     ax2.view_init(elev=90, azim=0)

    #     plt.show()